{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.8/site-packages/bitsandbytes/cextension.py:34: UserWarning: The installed version of bitsandbytes was compiled without GPU support. 8-bit optimizers, 8-bit multiplication, and GPU quantization are unavailable.\n",
      "  warn(\"The installed version of bitsandbytes was compiled without GPU support. \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.8/site-packages/bitsandbytes/libbitsandbytes_cpu.so: undefined symbol: cadam32bit_grad_fp32\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-07-27 15:03:53.011522: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-07-27 15:03:53.088161: I tensorflow/core/util/port.cc:104] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-07-27 15:03:53.512914: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory\n",
      "2023-07-27 15:03:53.512955: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory\n",
      "2023-07-27 15:03:53.512958: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      "/home/husein/.local/lib/python3.8/site-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.1.36ubuntu1 is an invalid version and will not be supported in a future release\n",
      "  warnings.warn(\n",
      "/home/husein/.local/lib/python3.8/site-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.23ubuntu1 is an invalid version and will not be supported in a future release\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    T5ForConditionalGeneration,\n",
    "    AutoConfig,\n",
    ")\n",
    "from pytorch_lightning import LightningModule\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('./out-tiny-1.1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.8/site-packages/transformers/utils/hub.py:665: UserWarning: The `organization` argument is deprecated and will be removed in v5 of Transformers. Set your organization directly in the `repo_id` passed instead (`repo_id={organization}/{model_id}`).\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/nanot5-tiny-malaysian-cased/commit/418eeb558f6a882d08c8b55bbed0444de73e2593', commit_message='Upload tokenizer', commit_description='', oid='418eeb558f6a882d08c8b55bbed0444de73e2593', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.push_to_hub('nanot5-tiny-malaysian-cased', organization='mesolitica')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Module(LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        config = AutoConfig.from_pretrained(\n",
    "            './out-tiny-1.1'\n",
    "        )\n",
    "        self.model = T5ForConditionalGeneration.from_pretrained(\n",
    "            './out-tiny-1.1',\n",
    "            config=config,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'model-epoch=00-step=120000.ckpt'  'model-epoch=00-step=450000.ckpt'\r\n",
      "'model-epoch=00-step=130000.ckpt'  'model-epoch=00-step=460000.ckpt'\r\n",
      "'model-epoch=00-step=140000.ckpt'  'model-epoch=00-step=470000.ckpt'\r\n"
     ]
    }
   ],
   "source": [
    "!ls logs/tiny"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Module()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = model.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_weights = torch.load('logs/tiny/model-epoch=00-step=470000.ckpt',\n",
    "                             map_location=torch.device('cpu'))['state_dict'].items()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model._orig_mod.shared.weight model.shared.weight\n",
      "model._orig_mod.encoder.embed_tokens.weight model.encoder.embed_tokens.weight\n",
      "model._orig_mod.encoder.block.0.layer.0.SelfAttention.q.weight model.encoder.block.0.layer.0.SelfAttention.q.weight\n",
      "model._orig_mod.encoder.block.0.layer.0.SelfAttention.k.weight model.encoder.block.0.layer.0.SelfAttention.k.weight\n",
      "model._orig_mod.encoder.block.0.layer.0.SelfAttention.v.weight model.encoder.block.0.layer.0.SelfAttention.v.weight\n",
      "model._orig_mod.encoder.block.0.layer.0.SelfAttention.o.weight model.encoder.block.0.layer.0.SelfAttention.o.weight\n",
      "model._orig_mod.encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight model.encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight\n",
      "model._orig_mod.encoder.block.0.layer.0.layer_norm.weight model.encoder.block.0.layer.0.layer_norm.weight\n",
      "model._orig_mod.encoder.block.0.layer.1.DenseReluDense.wi_0.weight model.encoder.block.0.layer.1.DenseReluDense.wi_0.weight\n",
      "model._orig_mod.encoder.block.0.layer.1.DenseReluDense.wi_1.weight model.encoder.block.0.layer.1.DenseReluDense.wi_1.weight\n",
      "model._orig_mod.encoder.block.0.layer.1.DenseReluDense.wo.weight model.encoder.block.0.layer.1.DenseReluDense.wo.weight\n",
      "model._orig_mod.encoder.block.0.layer.1.layer_norm.weight model.encoder.block.0.layer.1.layer_norm.weight\n",
      "model._orig_mod.encoder.block.1.layer.0.SelfAttention.q.weight model.encoder.block.1.layer.0.SelfAttention.q.weight\n",
      "model._orig_mod.encoder.block.1.layer.0.SelfAttention.k.weight model.encoder.block.1.layer.0.SelfAttention.k.weight\n",
      "model._orig_mod.encoder.block.1.layer.0.SelfAttention.v.weight model.encoder.block.1.layer.0.SelfAttention.v.weight\n",
      "model._orig_mod.encoder.block.1.layer.0.SelfAttention.o.weight model.encoder.block.1.layer.0.SelfAttention.o.weight\n",
      "model._orig_mod.encoder.block.1.layer.0.layer_norm.weight model.encoder.block.1.layer.0.layer_norm.weight\n",
      "model._orig_mod.encoder.block.1.layer.1.DenseReluDense.wi_0.weight model.encoder.block.1.layer.1.DenseReluDense.wi_0.weight\n",
      "model._orig_mod.encoder.block.1.layer.1.DenseReluDense.wi_1.weight model.encoder.block.1.layer.1.DenseReluDense.wi_1.weight\n",
      "model._orig_mod.encoder.block.1.layer.1.DenseReluDense.wo.weight model.encoder.block.1.layer.1.DenseReluDense.wo.weight\n",
      "model._orig_mod.encoder.block.1.layer.1.layer_norm.weight model.encoder.block.1.layer.1.layer_norm.weight\n",
      "model._orig_mod.encoder.block.2.layer.0.SelfAttention.q.weight model.encoder.block.2.layer.0.SelfAttention.q.weight\n",
      "model._orig_mod.encoder.block.2.layer.0.SelfAttention.k.weight model.encoder.block.2.layer.0.SelfAttention.k.weight\n",
      "model._orig_mod.encoder.block.2.layer.0.SelfAttention.v.weight model.encoder.block.2.layer.0.SelfAttention.v.weight\n",
      "model._orig_mod.encoder.block.2.layer.0.SelfAttention.o.weight model.encoder.block.2.layer.0.SelfAttention.o.weight\n",
      "model._orig_mod.encoder.block.2.layer.0.layer_norm.weight model.encoder.block.2.layer.0.layer_norm.weight\n",
      "model._orig_mod.encoder.block.2.layer.1.DenseReluDense.wi_0.weight model.encoder.block.2.layer.1.DenseReluDense.wi_0.weight\n",
      "model._orig_mod.encoder.block.2.layer.1.DenseReluDense.wi_1.weight model.encoder.block.2.layer.1.DenseReluDense.wi_1.weight\n",
      "model._orig_mod.encoder.block.2.layer.1.DenseReluDense.wo.weight model.encoder.block.2.layer.1.DenseReluDense.wo.weight\n",
      "model._orig_mod.encoder.block.2.layer.1.layer_norm.weight model.encoder.block.2.layer.1.layer_norm.weight\n",
      "model._orig_mod.encoder.block.3.layer.0.SelfAttention.q.weight model.encoder.block.3.layer.0.SelfAttention.q.weight\n",
      "model._orig_mod.encoder.block.3.layer.0.SelfAttention.k.weight model.encoder.block.3.layer.0.SelfAttention.k.weight\n",
      "model._orig_mod.encoder.block.3.layer.0.SelfAttention.v.weight model.encoder.block.3.layer.0.SelfAttention.v.weight\n",
      "model._orig_mod.encoder.block.3.layer.0.SelfAttention.o.weight model.encoder.block.3.layer.0.SelfAttention.o.weight\n",
      "model._orig_mod.encoder.block.3.layer.0.layer_norm.weight model.encoder.block.3.layer.0.layer_norm.weight\n",
      "model._orig_mod.encoder.block.3.layer.1.DenseReluDense.wi_0.weight model.encoder.block.3.layer.1.DenseReluDense.wi_0.weight\n",
      "model._orig_mod.encoder.block.3.layer.1.DenseReluDense.wi_1.weight model.encoder.block.3.layer.1.DenseReluDense.wi_1.weight\n",
      "model._orig_mod.encoder.block.3.layer.1.DenseReluDense.wo.weight model.encoder.block.3.layer.1.DenseReluDense.wo.weight\n",
      "model._orig_mod.encoder.block.3.layer.1.layer_norm.weight model.encoder.block.3.layer.1.layer_norm.weight\n",
      "model._orig_mod.encoder.final_layer_norm.weight model.encoder.final_layer_norm.weight\n",
      "model._orig_mod.decoder.embed_tokens.weight model.decoder.embed_tokens.weight\n",
      "model._orig_mod.decoder.block.0.layer.0.SelfAttention.q.weight model.decoder.block.0.layer.0.SelfAttention.q.weight\n",
      "model._orig_mod.decoder.block.0.layer.0.SelfAttention.k.weight model.decoder.block.0.layer.0.SelfAttention.k.weight\n",
      "model._orig_mod.decoder.block.0.layer.0.SelfAttention.v.weight model.decoder.block.0.layer.0.SelfAttention.v.weight\n",
      "model._orig_mod.decoder.block.0.layer.0.SelfAttention.o.weight model.decoder.block.0.layer.0.SelfAttention.o.weight\n",
      "model._orig_mod.decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight model.decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight\n",
      "model._orig_mod.decoder.block.0.layer.0.layer_norm.weight model.decoder.block.0.layer.0.layer_norm.weight\n",
      "model._orig_mod.decoder.block.0.layer.1.EncDecAttention.q.weight model.decoder.block.0.layer.1.EncDecAttention.q.weight\n",
      "model._orig_mod.decoder.block.0.layer.1.EncDecAttention.k.weight model.decoder.block.0.layer.1.EncDecAttention.k.weight\n",
      "model._orig_mod.decoder.block.0.layer.1.EncDecAttention.v.weight model.decoder.block.0.layer.1.EncDecAttention.v.weight\n",
      "model._orig_mod.decoder.block.0.layer.1.EncDecAttention.o.weight model.decoder.block.0.layer.1.EncDecAttention.o.weight\n",
      "model._orig_mod.decoder.block.0.layer.1.layer_norm.weight model.decoder.block.0.layer.1.layer_norm.weight\n",
      "model._orig_mod.decoder.block.0.layer.2.DenseReluDense.wi_0.weight model.decoder.block.0.layer.2.DenseReluDense.wi_0.weight\n",
      "model._orig_mod.decoder.block.0.layer.2.DenseReluDense.wi_1.weight model.decoder.block.0.layer.2.DenseReluDense.wi_1.weight\n",
      "model._orig_mod.decoder.block.0.layer.2.DenseReluDense.wo.weight model.decoder.block.0.layer.2.DenseReluDense.wo.weight\n",
      "model._orig_mod.decoder.block.0.layer.2.layer_norm.weight model.decoder.block.0.layer.2.layer_norm.weight\n",
      "model._orig_mod.decoder.block.1.layer.0.SelfAttention.q.weight model.decoder.block.1.layer.0.SelfAttention.q.weight\n",
      "model._orig_mod.decoder.block.1.layer.0.SelfAttention.k.weight model.decoder.block.1.layer.0.SelfAttention.k.weight\n",
      "model._orig_mod.decoder.block.1.layer.0.SelfAttention.v.weight model.decoder.block.1.layer.0.SelfAttention.v.weight\n",
      "model._orig_mod.decoder.block.1.layer.0.SelfAttention.o.weight model.decoder.block.1.layer.0.SelfAttention.o.weight\n",
      "model._orig_mod.decoder.block.1.layer.0.layer_norm.weight model.decoder.block.1.layer.0.layer_norm.weight\n",
      "model._orig_mod.decoder.block.1.layer.1.EncDecAttention.q.weight model.decoder.block.1.layer.1.EncDecAttention.q.weight\n",
      "model._orig_mod.decoder.block.1.layer.1.EncDecAttention.k.weight model.decoder.block.1.layer.1.EncDecAttention.k.weight\n",
      "model._orig_mod.decoder.block.1.layer.1.EncDecAttention.v.weight model.decoder.block.1.layer.1.EncDecAttention.v.weight\n",
      "model._orig_mod.decoder.block.1.layer.1.EncDecAttention.o.weight model.decoder.block.1.layer.1.EncDecAttention.o.weight\n",
      "model._orig_mod.decoder.block.1.layer.1.layer_norm.weight model.decoder.block.1.layer.1.layer_norm.weight\n",
      "model._orig_mod.decoder.block.1.layer.2.DenseReluDense.wi_0.weight model.decoder.block.1.layer.2.DenseReluDense.wi_0.weight\n",
      "model._orig_mod.decoder.block.1.layer.2.DenseReluDense.wi_1.weight model.decoder.block.1.layer.2.DenseReluDense.wi_1.weight\n",
      "model._orig_mod.decoder.block.1.layer.2.DenseReluDense.wo.weight model.decoder.block.1.layer.2.DenseReluDense.wo.weight\n",
      "model._orig_mod.decoder.block.1.layer.2.layer_norm.weight model.decoder.block.1.layer.2.layer_norm.weight\n",
      "model._orig_mod.decoder.block.2.layer.0.SelfAttention.q.weight model.decoder.block.2.layer.0.SelfAttention.q.weight\n",
      "model._orig_mod.decoder.block.2.layer.0.SelfAttention.k.weight model.decoder.block.2.layer.0.SelfAttention.k.weight\n",
      "model._orig_mod.decoder.block.2.layer.0.SelfAttention.v.weight model.decoder.block.2.layer.0.SelfAttention.v.weight\n",
      "model._orig_mod.decoder.block.2.layer.0.SelfAttention.o.weight model.decoder.block.2.layer.0.SelfAttention.o.weight\n",
      "model._orig_mod.decoder.block.2.layer.0.layer_norm.weight model.decoder.block.2.layer.0.layer_norm.weight\n",
      "model._orig_mod.decoder.block.2.layer.1.EncDecAttention.q.weight model.decoder.block.2.layer.1.EncDecAttention.q.weight\n",
      "model._orig_mod.decoder.block.2.layer.1.EncDecAttention.k.weight model.decoder.block.2.layer.1.EncDecAttention.k.weight\n",
      "model._orig_mod.decoder.block.2.layer.1.EncDecAttention.v.weight model.decoder.block.2.layer.1.EncDecAttention.v.weight\n",
      "model._orig_mod.decoder.block.2.layer.1.EncDecAttention.o.weight model.decoder.block.2.layer.1.EncDecAttention.o.weight\n",
      "model._orig_mod.decoder.block.2.layer.1.layer_norm.weight model.decoder.block.2.layer.1.layer_norm.weight\n",
      "model._orig_mod.decoder.block.2.layer.2.DenseReluDense.wi_0.weight model.decoder.block.2.layer.2.DenseReluDense.wi_0.weight\n",
      "model._orig_mod.decoder.block.2.layer.2.DenseReluDense.wi_1.weight model.decoder.block.2.layer.2.DenseReluDense.wi_1.weight\n",
      "model._orig_mod.decoder.block.2.layer.2.DenseReluDense.wo.weight model.decoder.block.2.layer.2.DenseReluDense.wo.weight\n",
      "model._orig_mod.decoder.block.2.layer.2.layer_norm.weight model.decoder.block.2.layer.2.layer_norm.weight\n",
      "model._orig_mod.decoder.block.3.layer.0.SelfAttention.q.weight model.decoder.block.3.layer.0.SelfAttention.q.weight\n",
      "model._orig_mod.decoder.block.3.layer.0.SelfAttention.k.weight model.decoder.block.3.layer.0.SelfAttention.k.weight\n",
      "model._orig_mod.decoder.block.3.layer.0.SelfAttention.v.weight model.decoder.block.3.layer.0.SelfAttention.v.weight\n",
      "model._orig_mod.decoder.block.3.layer.0.SelfAttention.o.weight model.decoder.block.3.layer.0.SelfAttention.o.weight\n",
      "model._orig_mod.decoder.block.3.layer.0.layer_norm.weight model.decoder.block.3.layer.0.layer_norm.weight\n",
      "model._orig_mod.decoder.block.3.layer.1.EncDecAttention.q.weight model.decoder.block.3.layer.1.EncDecAttention.q.weight\n",
      "model._orig_mod.decoder.block.3.layer.1.EncDecAttention.k.weight model.decoder.block.3.layer.1.EncDecAttention.k.weight\n",
      "model._orig_mod.decoder.block.3.layer.1.EncDecAttention.v.weight model.decoder.block.3.layer.1.EncDecAttention.v.weight\n",
      "model._orig_mod.decoder.block.3.layer.1.EncDecAttention.o.weight model.decoder.block.3.layer.1.EncDecAttention.o.weight\n",
      "model._orig_mod.decoder.block.3.layer.1.layer_norm.weight model.decoder.block.3.layer.1.layer_norm.weight\n",
      "model._orig_mod.decoder.block.3.layer.2.DenseReluDense.wi_0.weight model.decoder.block.3.layer.2.DenseReluDense.wi_0.weight\n",
      "model._orig_mod.decoder.block.3.layer.2.DenseReluDense.wi_1.weight model.decoder.block.3.layer.2.DenseReluDense.wi_1.weight\n",
      "model._orig_mod.decoder.block.3.layer.2.DenseReluDense.wo.weight model.decoder.block.3.layer.2.DenseReluDense.wo.weight\n",
      "model._orig_mod.decoder.block.3.layer.2.layer_norm.weight model.decoder.block.3.layer.2.layer_norm.weight\n",
      "model._orig_mod.decoder.final_layer_norm.weight model.decoder.final_layer_norm.weight\n",
      "model._orig_mod.lm_head.weight model.lm_head.weight\n"
     ]
    }
   ],
   "source": [
    "for k, v in old_weights:\n",
    "    new_k = k.replace('._orig_mod', '')\n",
    "    print(k, new_k)\n",
    "    weights[new_k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "968edf82baaa4e03aad58a7896a66c32",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/205M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/nanot5-tiny-malaysian-cased/commit/37d0092ffcd17e9112ab9d6eeb95a6c9b47520b2', commit_message='Upload T5ForConditionalGeneration', commit_description='', oid='37d0092ffcd17e9112ab9d6eeb95a6c9b47520b2', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.model.push_to_hub('nanot5-tiny-malaysian-cased', organization='mesolitica')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
